package com.jingding.procurement.unit;
import com.jingding.procurement.ex.ServiceException;
import com.jingding.procurement.restful.ServiceCode;
import org.springframework.stereotype.Component;
import org.springframework.web.servlet.HandlerInterceptor;

import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;
import java.util.HashMap;
import java.util.Map;
import java.util.Timer;
import java.util.TimerTask;


@Component
public class RateLimitInterceptor implements HandlerInterceptor {

    private static final int MAX_REQUESTS = 10; // 同一时间段内允许的最大请求数
    private static final long TIME_PERIOD = 600 * 1000; // 时间段，单位为毫秒 在一分钟内限制ip访问次数为10次

    private Map<String, Integer> requestCounts = new HashMap<String, Integer>();

    @Override
    public boolean preHandle(HttpServletRequest request, HttpServletResponse response, Object handler) throws Exception {
        final String ipAddress = request.getRemoteAddr();
        System.out.println(ipAddress);
        // 检查 IP 地址是否已经达到最大请求数
        if (requestCounts.containsKey(ipAddress) && requestCounts.get(ipAddress) >= MAX_REQUESTS) {
            throw new ServiceException(ServiceCode.IP_LIMIT.getCode(),ServiceCode.IP_LIMIT.getMessage());
        }

        // 更新 IP 地址的请求数
        requestCounts.put(ipAddress, requestCounts.getOrDefault(ipAddress, 0) + 1);

        // 在指定时间后清除 IP 地址的请求数
        new Timer().schedule(
                new TimerTask() {
                    @Override
                    public void run() {
                        requestCounts.remove(ipAddress);
                    }
                },
                TIME_PERIOD
        );
        return true;
    }
}
